#ifndef HERO_VISION__PNP_SOLVER_HPP_
#define HERO_VISION__PNP_SOLVER_HPP_

#include <geometry_msgs/msg/point.hpp>
#include <opencv2/core.hpp>

// STD
#include <array>
#include <vector>

#include "hero_integrated_vision/armor.hpp"
#include "aim_interfaces/msg/armor.hpp"

namespace rm_aim
{
class PnPSolver
{
public:
  PnPSolver(
    const std::array<double, 9> & camera_matrix,
    const std::vector<double> & distortion_coefficients);

  // Get 3d position
  bool solvePnP(const Armor & armor, cv::Mat & rvec, cv::Mat & tvec);

  // Calculate the distance between armor center and image center
  float calculateDistanceToCenter(const cv::Point2f & image_point);

  /**
   * @brief solvePnP后直接优化yaw，并返回完整Pose（含优化后的orientation）
   * @param rvec 原始旋转向量（solvePnP输出）
   * @param tvec 原始平移向量（solvePnP输出，位置不变）
   * @param armor 装甲板检测信息
   * @param camera_matrix 相机内参
   * @param distort_coeffs 畸变系数
   * @return 优化后的Pose（position=原始tvec，orientation=优化yaw后的四元数）
   */
  geometry_msgs::msg::Pose optimize_yaw_direct(
    const cv::Mat& rvec,
    const cv::Mat& tvec,
    const Armor& armor,
    const cv::Mat& camera_matrix,
    const cv::Mat& distort_coeffs) const;

private:
  cv::Mat camera_matrix_;
  cv::Mat dist_coeffs_;

  // Unit: mm
  static constexpr float SMALL_ARMOR_WIDTH = 135;
  static constexpr float SMALL_ARMOR_HEIGHT = 55;
  static constexpr float LARGE_ARMOR_WIDTH = 225;
  static constexpr float LARGE_ARMOR_HEIGHT = 55;

  // Four vertices of armor in 3d
  std::vector<cv::Point3f> small_armor_points_;
  std::vector<cv::Point3f> large_armor_points_;

  /**
   * @brief 直接计算重投影误差
   * @param yaw_candidate 候选yaw
   * @param pitch 原始pitch
   * @param roll 原始roll
   * @param tvec 平移向量
   * @param armor 装甲板信息
   * @param camera_matrix 相机内参
   * @param distort_coeffs 畸变系数
   * @return 总重投影误差
   */
  double calculate_reprojection_error_direct(
    double yaw_candidate,
    double pitch,
    double roll,
    const cv::Mat& tvec,
    const Armor& armor,
    const cv::Mat& camera_matrix,
    const cv::Mat& distort_coeffs) const;
};

}  // namespace rm_aim

#endif  // HERO_VISION__PNP_SOLVER_HPP_
